#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/7/29 10:57
# @Author  : YueJian
# @File    : captcha.py
# @Description : 图片验证码生成工具类
import base64
import os.path
import random
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from io import BytesIO
from typing import Tuple, Optional, Union
from src.config import settings
from PIL import Image, ImageDraw, ImageFont, ImageFilter


class CaptchaType(Enum):
    """验证码类型"""

    TEXT = "text"  # 文本验证码
    ARITHMETIC = "arithmetic"  # 算术验证码


@dataclass
class CaptchaConfig:
    """验证码配置类"""

    # 通用配置
    captcha_type: CaptchaType = CaptchaType.TEXT
    width: int = 160  # 统一宽度
    height: int = 60  # 统一高度
    font_path: str = os.path.join(
        settings.system.BASE_PATH, "resources", "font", "font.ttf"
    )
    font_size: int = 24  # 统一字体大小

    # 文本验证码配置
    length: int = 4
    use_curve: bool = True
    use_noise: bool = True

    # 算术验证码配置


class CaptchaStrategy(ABC):
    """验证码生成策略基类"""

    def __init__(self, config: CaptchaConfig):
        self.config = config
        self._draw: Optional[ImageDraw.Draw] = None
        self._image: Optional[Image.Image] = None
        self._code: Optional[Union[list[str], int]] = None
        self._char_set: Optional[str] = None

    def _create_image(self) -> None:
        """创建画布"""
        bg_color = (
            random.randint(200, 255),
            random.randint(200, 255),
            random.randint(200, 255),
        )
        self._image = Image.new(
            "RGB", (self.config.width, self.config.height), bg_color
        )
        self._draw = ImageDraw.Draw(self._image)

    def _add_noise_lines(self) -> None:
        """添加干扰线"""
        if not self._draw:
            return

        # 添加3-5条干扰线
        for _ in range(random.randint(3, 5)):
            # 随机生成线条的起点和终点
            start_x = random.randint(0, self.config.width)
            start_y = random.randint(0, self.config.height)
            end_x = random.randint(0, self.config.width)
            end_y = random.randint(0, self.config.height)

            # 随机生成线条颜色（使用较浅的颜色）
            line_color = (
                random.randint(150, 200),
                random.randint(150, 200),
                random.randint(150, 200),
            )

            # 随机生成线条宽度
            line_width = random.randint(1, 2)

            # 绘制线条
            self._draw.line(
                [(start_x, start_y), (end_x, end_y)], fill=line_color, width=line_width
            )

    def _apply_filters(self) -> None:
        """应用滤镜效果"""
        if not self._image:
            return

        # 移除模糊效果，只保留轻微的模式滤镜
        self._image = self._image.filter(
            ImageFilter.ModeFilter(1)
        )  # 降低模式滤镜的强度

    def _draw_code(self) -> None:
        """绘制验证码"""
        if not self._draw or not self._code:
            return

        font = ImageFont.truetype(self.config.font_path, self.config.font_size)
        chars = self._get_chars_to_draw()

        for i, char in enumerate(chars):
            # 使用更深的颜色
            color = (
                random.randint(0, 100),  # 降低亮度范围
                random.randint(0, 100),
                random.randint(0, 100),
            )
            # 计算字符位置
            x, y = self._calculate_char_position(i, len(chars))
            self._draw.text((x, y), char, font=font, fill=color)

    def _calculate_char_position(
        self, index: int, total_chars: int
    ) -> Tuple[float, float]:
        """计算字符位置

        Args:
            index: 字符索引
            total_chars: 总字符数

        Returns:
            Tuple[float, float]: 字符的x和y坐标
        """
        x = self.config.font_size * index + random.uniform(2, 4)
        y = random.uniform(2, 4)
        return x, y

    @abstractmethod
    def _get_chars_to_draw(self) -> list[str]:
        """获取要绘制的字符列表"""
        pass

    @abstractmethod
    def _generate_code(self) -> None:
        """生成验证码字符"""
        pass

    @abstractmethod
    def generate(
        self, size: Optional[int] = None, length: Optional[int] = None
    ) -> Tuple[Image.Image, Union[str, int]]:
        """生成验证码"""
        pass


class TextCaptchaStrategy(CaptchaStrategy):
    """文本验证码策略"""

    def _generate_code(self) -> None:
        """生成验证码字符"""
        if not self._char_set:
            number = "3456789"
            letters = "qwertyuipasdfghjkzxcvbnmQWERTYUIOPASDFGHJKLZXCVBNM"
            self._char_set = letters + number

        self._code = random.sample(self._char_set, self.config.length)

    def _get_chars_to_draw(self) -> list[str]:
        """获取要绘制的字符列表"""
        return self._code if self._code else []

    def _add_noise(self) -> None:
        """添加干扰点"""
        if not self.config.use_noise or not self._draw:
            return

        font = ImageFont.truetype(
            self.config.font_path, int(self.config.font_size / 2)
        )  # 减小干扰点大小
        for _ in range(5):  # 减少干扰点数量
            noise_color = (
                random.randint(150, 200),
                random.randint(150, 200),
                random.randint(150, 200),
            )
            noise_chars = random.sample(self._char_set, 3)  # 减少每个干扰点的字符数
            for char in noise_chars:
                pos = (
                    random.randint(0, self.config.width),
                    random.randint(0, self.config.height),
                )
                self._draw.text(pos, char, font=font, fill=noise_color)

    def generate(
        self, size: Optional[int] = None, length: Optional[int] = None
    ) -> Tuple[Image.Image, str]:
        """生成文本验证码"""
        if size:
            self.config.font_size = size
        if length:
            self.config.length = length

        self.config.width = self.config.font_size * self.config.length
        self.config.height = int(self.config.font_size * 1.5)

        self._create_image()
        self._generate_code()
        self._add_noise()
        self._draw_code()
        self._apply_filters()

        return self._image, "".join(self._code).lower()


class ArithmeticCaptchaStrategy(CaptchaStrategy):
    """算术验证码策略"""

    def _generate_code(self) -> None:
        """生成验证码字符"""
        num1 = random.randint(0, 9)
        num2 = random.randint(0, 9)
        operator = random.choice(["+", "-", "*"])

        if operator == "+":
            result = num1 + num2
        elif operator == "-":
            result = num1 - num2
        else:
            result = num1 * num2

        self._code = [num1, operator, num2, result]

    def _get_chars_to_draw(self) -> list[str]:
        """获取要绘制的字符列表"""
        if not self._code:
            return []
        return list(f"{self._code[0]} {self._code[1]} {self._code[2]} = ?")

    def _calculate_char_position(
        self, index: int, total_chars: int
    ) -> Tuple[float, float]:
        """计算字符位置

        Args:
            index: 字符索引
            total_chars: 总字符数

        Returns:
            Tuple[float, float]: 字符的x和y坐标
        """
        x = self.config.font_size * index * 0.6 + random.uniform(2, 4)
        y = random.uniform(2, 4)
        return x, y

    def generate(
        self, size: Optional[int] = None, length: Optional[int] = None
    ) -> Tuple[Image.Image, str]:
        """生成算术验证码"""
        if size:
            self.config.font_size = size

        # 调整宽度计算方式，使其更紧凑
        # 算术表达式格式为 "X + Y = ?"，总共7个字符（包括空格）
        # 每个字符宽度约为字体大小的0.6倍，加上一些边距
        char_width = self.config.font_size * 0.7
        self.config.width = int(char_width * 7 + 20)  # 7个字符加上一些边距
        self.config.height = int(self.config.font_size * 1.5)

        self._create_image()
        self._generate_code()
        self._add_noise_lines()
        self._draw_code()
        self._apply_filters()

        return self._image, str(self._code[3])


class CaptchaGenerator:
    """验证码生成器"""

    def __init__(self, config: Optional[CaptchaConfig] = None):
        self.config = config or CaptchaConfig()
        self._strategy = self._create_strategy()

    def _create_strategy(self) -> CaptchaStrategy:
        """创建验证码生成策略"""
        if self.config.captcha_type == CaptchaType.TEXT:
            return TextCaptchaStrategy(self.config)
        else:
            return ArithmeticCaptchaStrategy(self.config)

    def generate(
        self, size: Optional[int] = None, length: Optional[int] = None
    ) -> Tuple[Image.Image, Union[str, int]]:
        """生成验证码图片和验证码"""
        return self._strategy.generate(size, length)


def get_captcha(
    captcha_type: CaptchaType = CaptchaType.TEXT,
) -> Tuple[str, Union[str, int]]:
    """获取验证码图片和验证码

    Args:
        captcha_type: 验证码类型，支持文本验证码和算术验证码

    Returns:
        Tuple[str, Union[str, int]]: base64编码的图片和验证码（文本验证码返回字符串，算术验证码返回整数）
    """
    config = CaptchaConfig(captcha_type=captcha_type)
    generator = CaptchaGenerator(config)
    image, code = generator.generate()

    # 将图像保存到内存缓冲区
    buffer = BytesIO()
    image.save(buffer, format="JPEG")
    buffer.seek(0)

    # 转换为base64
    b64_img = (
        f"data:image/jpeg;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
    )
    return b64_img, code
